# Pattern matching
# Rule matching
import json
from re import L
from miasm.core.asmblock import AsmCFG
from miasm.expression.expression import ExprLoc, ExprInt, ExprOp, ExprAssign,ExprCond, ExprId, ExprMem,ExprSlice
# 特征
# 得到平坦度
def get_flatness():
    
    return




# match_seq：匹配的指令序列
# seq_num:指令序列的匹配数
# blk_line:指令行
# store_obj:保存r/m的记录
def match_line(match_seq,seq_num,blk_line,store_obj):
    op_name=blk_line.name
    op_args=blk_line.args
    match_single_rule=match_seq[seq_num]
    match_op_name=match_single_rule['op_name']
    match_operand=match_single_rule['operand']
    # check op name & operand num
    if (match_op_name!=op_name) or (len(op_args)!=len(match_operand)):
        return False
    # check operand type
    for op_num in range(0,len(op_args)):
        rule_arg=match_operand[op_num]
        real_arg=op_args[op_num]
        # check arg type
        # str
        if isinstance(rule_arg,str):
            if (rule_arg[0]=='r' and isinstance(real_arg,ExprId)) \
                or (rule_arg[0]=='m' and isinstance(real_arg,ExprMem)):
                # check content
                content_index=int(rule_arg[1])
                if rule_arg[0]=='r':
                    if store_obj['reg'][content_index] is None:
                        store_obj['reg'][content_index]=real_arg
                    elif store_obj['reg'][content_index]==real_arg:
                        pass
                    else:
                        return False
                elif rule_arg[0]=='m':
                    if store_obj['mem'][content_index] is None:
                        store_obj['mem'][content_index]=real_arg
                    elif store_obj['mem'][content_index]==real_arg:
                        pass
                    else:
                        return False
                else:
                    assert()
                pass
            else:
                return False
        # imm
        elif isinstance(rule_arg,int):
            if isinstance(real_arg,ExprInt) and rule_arg==int(real_arg):
                continue
            return False    
        else:
            print(rule_arg)
            return False
    #print('[debug]match',match_single_rule)
    #print(seq_num,op_name,op_args)
    return True

# 判断asmblk从第line_num开始是否符合match_seq规则
# 匹配指令序列 match_seq:匹配的序列,asmblk:汇编指令块,line_num：匹配的块行号
def do_match_seq(match_seq,asmblk,line_num):
    target_num=line_num
    store_obj={
        'reg':[None]*8,
        'mem':[None]*8
    }
    # 序列长度
    seq_len=len(match_seq)
    instrs_len=len(asmblk.lines)-line_num
    if seq_len > instrs_len:
        return False,None
    # 汇编块
    for seq_num in range(0,seq_len):
        # 匹配的序列
        match_res=match_line(match_seq,seq_num,asmblk.lines[line_num],store_obj)
        if not match_res:
            return False,None
        line_num+=1
    
    return True,store_obj

# 替换指定的指令
def replace_instrs(match_seq,fix_seq,asmblk,line_num,store_obj):
    match_instrs_num=len(match_seq)
    fix_instrs_num=len(fix_seq)
    if fix_instrs_num>match_instrs_num:
        # 修补后指令的数量大于修补前指令数量
        assert()
    nop_num=match_instrs_num-fix_instrs_num
    
    for n in range(0,match_instrs_num):
        fix_line_num=line_num+n
        if n<fix_instrs_num:
            fix_op_name=fix_seq[n]['op_name']
            fix_args=fix_seq[n]['operand']
            asmblk.lines[fix_line_num].name=fix_op_name
            # clear
            asmblk.lines[fix_line_num].args.clear()
            for fix_arg in fix_args:
                if isinstance(fix_arg,str):
                    content_index=int(fix_arg[1])
                    if fix_arg[0]=='r':
                        asmblk.lines[fix_line_num].args.append(store_obj['reg'][content_index])
                    elif fix_arg[0]=='m':
                        asmblk.lines[fix_line_num].args.append(store_obj['mem'][content_index])
                    else:
                        assert()
                elif isinstance(fix_arg,int):
                    asmblk.lines[fix_line_num].args.append(ExprInt(fix_arg,32))
                else:
                    assert()
        else:
            asmblk.lines[fix_line_num].args.clear()
            asmblk.lines[fix_line_num].name='NOP'
    
    print(asmblk)                
    
    return
def pattern_matching(asmcfg:AsmCFG):
    # 读取replace_config.json
    with open ('D:/code/idapython/project/replace_config.json','r') as f:
        f_data=json.load(f)
        print(f_data)
    rule_list=f_data['replace_list']
    for asmblk in asmcfg.blocks:
        for rule_item in rule_list:
            # 要匹配的规则
            match_seq=rule_item['match_sequence']
            match_num=0
            # 匹配每一行
            for line_num in range(0,len(asmblk.lines)):
                match_seq_ret,match_seq_obj=do_match_seq(match_seq,asmblk,line_num)
                # 对匹配规则进行指令替换
                if match_seq_ret:
                    fix_seq=rule_item['fix_sequence']   # 修补指令
                    replace_instrs(match_seq,fix_seq,asmblk,line_num,match_seq_obj)
                    print('match rules',rule_item,line_num)
    
    return

def rule_matching():
    
    return

if __name__=='__main__':
    assert()
    pattern_matching(1)
